package com.lili.collection.jdbc.plugins;

import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.StrUtil;
import com.lili.collection.core.context.BatchRecordContext;
import com.lili.collection.core.context.StageContext;
import com.lili.collection.core.context.impl.GenericBatchRecordContext;
import com.lili.collection.core.error.InputRunnerException;
import com.lili.collection.core.error.StageRecordException;
import com.lili.collection.core.fields.Field;
import com.lili.collection.core.plugins.InputPlugin;
import com.lili.collection.core.record.impl.CommonRecord;
import com.lili.collection.core.runner.result.ErrOutputData;
import com.lili.collection.core.runner.result.RunnerResult;
import com.lili.collection.jdbc.config.JdbcInputConfig;
import com.lili.collection.jdbc.enums.DatabaseVendor;
import com.lili.collection.jdbc.enums.UnknownTypeAction;
import com.lili.collection.jdbc.utils.JdbcUtil;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * JdbcInputPlugin
 *
 * @author lili
 * @date 2023/2/13 20:43
 */
@Slf4j
public class JdbcInputPlugin extends JdbcPlugin<JdbcInputConfig> implements InputPlugin<JdbcInputConfig> {

    //批量索引
    int batchIndex = -1;
    //下阶段偏移量值
    private String nextSourceOffset = null;
    //数据库连接
    private Connection connection = null;
    //返回结果集
    private ResultSet resultSet = null;
    //成功的总行
    int successfulTotalRowCount = 0;
    int totalRowCount = 0;
    int totalNumQueryErrors = 0;
    private boolean isFirstRun = true;
    //别名配置
    private Map<String, String> aliasMap;

    @Override
    public StageContext execute() {
        BatchRecordContext recordContext = new GenericBatchRecordContext();
        recordContext.setIndex(++batchIndex);
        this.executeNum++;

        int rowCount = 0;
        int numQueryErrors = 0;
        String sql = prepareQuery();
        int batchSize = config.getBatchSize();
        boolean firstBatchError=false;
        try {
            PreparedStatement ps = null;
            if (!(resultSet != null && resultSet.isClosed() && !isFirstRun)) {
                if (resultSet == null || resultSet.isClosed()) {
                    successfulTotalRowCount = 0;
                    int fetchSize = batchSize;
                    if (databaseVendor.isOneOf(DatabaseVendor.MYSQL)) {
                        // 配合TYPE_FORWARD_ONLY启动流模式
                        fetchSize = Integer.MIN_VALUE;
                    }
                    JdbcUtil.closeQuietly(connection);
                    connection = getProduceConnection();
                    totalRowCount = this.getTableTotal(connection, sql);
                    ps = connection.prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
                    ps.setFetchSize(fetchSize);
                    resultSet = ps.executeQuery();
                }
                String nextOffset = getNextSourceOffset();
                log.info("-->第{}批次数据开始读取，当前偏移量：{}", recordContext.getIndex(), nextOffset);
                boolean haveNext = true;
                while (continueReading(rowCount, batchSize) && (haveNext = resultSet.next())) {
                    LinkedHashMap<String, Field> fields = JdbcUtil.resultSetToFields(
                            resultSet,
                            config.getMaxClobSize(),
                            config.getMaxBlobSize(),
                            Collections.emptyMap(),
                            UnknownTypeAction.THROW_EXCEPTION,
                            aliasMap.keySet(),
                            false,
                            databaseVendor,
                            aliasMap
                    );
                    Field row = Field.create(fields);
                    if (config.isIncrementalMode()) {
                        nextSourceOffset = resultSet.getString(config.getOffsetColumn());
                    }
                    ++rowCount;
                    ++successfulTotalRowCount;
                    recordContext.setRecord(new CommonRecord(row));
                }
            }
            this.isFirstRun = false;
        } catch (Exception e) {
            if(executeNum==1){
                firstBatchError=true;
            }
            ++totalNumQueryErrors;
            recordContext.setException(new StageRecordException(e.getLocalizedMessage(), e));
            recordContext.setStageRecordState(StageContext.StageRecordState.ALL_ERROR);
            log.info("-->第{}批次数据读取时发生异常，原因：{}\r\n-->成功读取：{}，异常读取：{}，最新偏移量：{}",
                    e.getLocalizedMessage(),
                    recordContext.getIndex(),
                    rowCount,
                    numQueryErrors,
                    nextSourceOffset);
        }
        RunnerResult runnerResult = this.getRunnerResult();
        if (recordContext.getException() == null && recordContext.getOutput().size() == 0) {
            recordContext.setStageContextState(StageContext.StageContextState.COMPLETE);
            this.isFirstRun = true;
            runnerResult.setTotalRecordCount(totalRowCount);
            runnerResult.setSuccessRecordCount(successfulTotalRowCount);
            log.info("-->数据全部读取完成，共执行批次：{}\r\n-->总数据：{}，异常数：{}，最新偏移量：{}",
                    recordContext.getIndex(),
                    successfulTotalRowCount,
                    totalNumQueryErrors,
                    nextSourceOffset);
        } else {
            if (recordContext.getException() != null) {
                runnerResult.setErrorRecordCount(totalNumQueryErrors);
                runnerResult.setTotalRecordCount(totalRowCount);
                ErrOutputData errOutputData = runnerResult.getErrOutputData();
                errOutputData.setType("INPUT");
                errOutputData.setSql(sql);
                errOutputData.setTotalSize(totalRowCount);
                errOutputData.setBatchSize(batchSize);
                errOutputData.setBatchIndex(recordContext.getIndex());
                errOutputData.setFirstBatchError(firstBatchError);
                errOutputData.setErrorMessage(recordContext.getException().toString());
                runnerResult.setSuccessed(false);
                if(config.isIncrementalMode()){
                    //报错且 增量/追加模式
                    runnerResult.setSuccessRecordCount(successfulTotalRowCount - rowCount);
                }else {
                    //报错且全量模式
                    runnerResult.setSuccessRecordCount(0);
                }
            }
            recordContext.setStageContextState(StageContext.StageContextState.CONTINUE);
            log.info("-->第{}批次数据读取完成\r\n-->成功读取：{}，异常读取：{}，最新偏移量：{}",
                    recordContext.getIndex(),
                    rowCount,
                    numQueryErrors,
                    nextSourceOffset);
        }
        return recordContext;
    }

    @Override
    public String getNextSourceOffset() {
        if (StrUtil.isNotEmpty(nextSourceOffset)) {
            return nextSourceOffset;
        }
        if (config.isIncrementalMode()) {
            String lastOffset = null;
            String nextSourceOffset = StrUtil.isEmpty(lastOffset) ? config.getInitialOffset() : lastOffset;
            return nextSourceOffset;
        }
        return config.getInitialOffset();
    }

    @Override
    public void dispose() {
        JdbcUtil.closeQuietly(this.dataSource);
    }

    @SneakyThrows
    private Connection getProduceConnection() {
        if (connection == null || connection.isClosed()) {
            connection = dataSource.getConnection();
            return connection;
        }
        return connection;
    }

    private boolean continueReading(int rowCount, int batchSize) {
        if (config.isIncrementalMode()) {
            return rowCount < batchSize;
        } else if (config.isNextMode()) {
            return rowCount < batchSize;
        } else {
            return true;
        }
    }

    @Override
    public void init(JdbcInputConfig config) {
        try {
            super.init(config);
            this.aliasMap = config.getAliasMap();
            if (MapUtil.isEmpty(this.aliasMap)) {
                throw new IllegalArgumentException("JDBC输入映射配置为空！");
            }
        } catch (Throwable e) {
            log.error("初始化异常!",e);
            throw new InputRunnerException(e);
        }
    }

    private String prepareQuery() {
        if (config.isIncrementalMode()) {
            final String offset = getNextSourceOffset();
            return config.getSql().replaceAll("\\$\\{(offset|OFFSET)}", offset);
        } else return config.getSql();
    }

    private int getTableTotal(Connection connection, String querySql) {
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            String sql = "select count(1) total from (" + querySql + ") tb";
            statement = connection.createStatement();
            resultSet = statement.executeQuery(sql);
            if (resultSet.next()) {
                return resultSet.getInt("total");
            }
        } catch (Exception e) {
            log.error("获取表总数失败",e);
            throw new InputRunnerException(e);
        } finally {
            JdbcUtil.closeDb(resultSet, statement, null);
        }
        return -1;
    }
}
